# 写入lib.rs文件
# test.rs文件中写入use
# 测试文件中重新封装test函数
import json
import os
import logging

from extract_component import extract_extern_declaration, extract_func_signature_from_extern_declaration, extract_function_names
from extract_component import extract_struct_declaration, extract_type_declaration, extract_c_label_blocks, extract_static_declaration
from process_test_functions import comment_code
from split_function import split_rust_functions_in_file
from global_config import CRATE_NAME, COMPARE_HASH_SRC_LIST


def get_all_module(path):
    PROJ_DIR = path
    SRC_DIR = os.path.join(PROJ_DIR, "src")
    modules = []
    for root, dirs, files in os.walk(SRC_DIR):
        for file in files:
            if "lib" in file or "alloc_testing" in file:
                continue
            modules.append(file.split('.')[0])
    return modules

def add_crate_and_lib(path):
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    SRC_DIR = os.path.join(PROJ_DIR, "src")
    LIB_FILE = os.path.join(SRC_DIR, "lib.rs")
    # 文件路径
    # 清空或创建 lib.rs 文件
    open(LIB_FILE, 'w').close()
    # 遍历指定目录下的所有文件
    for filename in os.listdir(SRC_DIR):
        file_path = os.path.join(SRC_DIR, filename)
        # 确保是文件而非目录
        if os.path.isfile(file_path) and filename != "lib.rs":
            # 去除文件后缀，只保留文件名
            modulename = os.path.splitext(filename)[0]
            # 将 pub mod 文件名; 写入 lib.rs 文件
            with open(LIB_FILE, 'a') as lib_file:
                lib_file.write(f"pub mod {modulename};\n")
    logging.info("lib.rs 文件已更新。")


    logging.info(f"add import for test file")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            with open(file_path, 'r') as f:
                lines = f.readlines()
            for i, line in enumerate(lines):
                if line.strip() == '' or line.startswith("#"):
                    continue

                if "cpp" in file:
                    # print("hhh")
                    modules = get_all_module(path)
                    # print(modules)
                    for m in modules:
                        lines.insert(i, f"use {CRATE_NAME}::{m}::*;\n")
                else:
                    lines.insert(i, f"use {CRATE_NAME}::hash_int::*;\n")
                    lines.insert(i, f"use {CRATE_NAME}::hash_string::*;\n")
                    lines.insert(i, f"use {CRATE_NAME}::hash_pointer::*;\n")
                    lines.insert(i, f"use {CRATE_NAME}::compare_int::*;\n")
                    lines.insert(i, f"use {CRATE_NAME}::compare_string::*;\n")
                    lines.insert(i, f"use {CRATE_NAME}::compare_pointer::*;\n")
                    if "compare_functions" not in file and "hash_functions" not in file:
                        lines.insert(i, f"use {CRATE_NAME}::{file.removesuffix('.rs').removeprefix('test_')}::*;\n")
                
                lines.insert(i, f"use {CRATE_NAME}::alloc_testing::*;\n")

                lines.insert(i, "use std::cell::RefCell;\n")
                lines.insert(i, "use std::sync::atomic::{AtomicUsize, Ordering};\n")
                lines.insert(i, "use ctor::dtor;\n")
                lines.insert(i, "extern crate libc;\n")
                break
            with open(file_path, 'w') as f:
                f.writelines(lines)

    ignore=["compare_int.rs","compare_string.rs","compare_pointer.rs",
        "hash_int.rs","hash_string.rs","hash_pointer.rs"]
    for root, dirs, files in os.walk(SRC_DIR):
        for file in files:
            logging.info(f"add import for {file}")
            # print(f"add import for {file}")
            if  "lib.rs" in file or file in ignore:
                continue
            
            file_path = os.path.join(root, file)
            with open(file_path, 'r') as f:
                lines = f.readlines()

            
            
            
            for i, line in enumerate(lines):
                if line.strip() == '' or line.startswith("#"):
                    continue
                if "alloc_testing.rs" in file:
                    lines.insert(i, "use std::cell::RefCell;\n")

                    lines.insert(i, "use std::mem;\n")
                else:
                    lines.insert(i, "use crate::alloc_testing::*;\n")
                    lines.insert(i, "use std::mem;\n")
                break
            with open(file_path, 'w') as f:
                f.writelines(lines)


def add_use_crate(proj_path, content):
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")

    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)


            with open(file_path, 'r') as f:
                lines = f.readlines()
            for i, line in enumerate(lines):
                if line.strip() == '' or line.startswith("#"):
                    continue
                lines.insert(i, content)
                break
            with open(file_path, 'w') as f:
                f.writelines(lines)


# 为source文件移除extern中的所有alloc_test函数
def remove_extern_declaration_for_srcfile(proj_path):
    PROJ_DIR = proj_path
    SRC_DIR = os.path.join(PROJ_DIR, "src")

    for root, dirs, files in os.walk(SRC_DIR):
        for file in files:
            file_path = os.path.join(root, file)

            res = extract_extern_declaration(file_path)
            extern_decl = res if res else ""

            func_name_to_sig = extract_func_signature_from_extern_declaration(extern_decl)
            if func_name_to_sig:
                for func_name in func_name_to_sig.keys():
                    if "alloc_test" in func_name:
                        comment_code(file_path=file_path, target_code=func_name_to_sig[func_name])

def remove_extern_declaration_for_testfile(proj_path):    
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    SRC_DIR = os.path.join(PROJ_DIR, "src")

    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)

            # 文件名 test_avl_tree.rs
            file_name = os.path.basename(file_path)
            # 提取模块名 avl_tree
            start_index = len("test_")
            end_index = file_name.find(".rs")
            module_name = file_name[start_index:end_index]

            remove_list = set()
            for file in COMPARE_HASH_SRC_LIST:
                res = extract_function_names(file)
                if res:
                    remove_list = remove_list | set(res)
            if "hash_functions" not in module_name and "compare_functions" not in module_name:
                if "cpp" in module_name:
                    modules = get_all_module(proj_path)
                    for m in modules:
                        res = extract_function_names(f"{SRC_DIR}/{m}.rs")
                        if res:
                            remove_list = remove_list | set(res)
                else:
                    res = extract_function_names(f"{SRC_DIR}/{module_name}.rs")
                    if res:
                        remove_list = remove_list | set(res)
    
            res = extract_extern_declaration(file_path)
            extern_decl = res if res else ""
            func_name_to_sig = extract_func_signature_from_extern_declaration(extern_decl)
            if func_name_to_sig:
                for func_name in func_name_to_sig.keys():
                    if func_name in remove_list or "run_tests" in func_name or "alloc_test" in func_name:
                        comment_code(file_path=file_path, target_code=func_name_to_sig[func_name])
    
def remove_struct_and_type_declaration(proj_path):
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    SRC_DIR = os.path.join(PROJ_DIR, "src")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            test_file_path = os.path.join(root, file)
            start_index = len("test_")
            src_file = file[start_index:]
            src_file_path = os.path.join(SRC_DIR, src_file)

            res = extract_struct_declaration(test_file_path)
            test_struct_codes = res if res else []

            res = extract_struct_declaration(src_file_path)
            src_struct_codes = res if res else []
            # logging.info(f"查找到{len(struct_codes)}个结构体定义")
            for struct in test_struct_codes:
                if struct in src_struct_codes:
                    comment_code(test_file_path, struct)
            attribute="#[derive(Copy, Clone)]\n#[repr(C)]"
            comment_code(test_file_path, attribute)

            res = extract_type_declaration(test_file_path)
            test_type_codes = res if res else []
            res = extract_type_declaration(src_file_path)
            src_type_codes = res if res else []
            # logging.info(f"查找到{len(type_codes)}个type定义")
            for type in test_type_codes:
                if "cpp" in file:
                    comment_code(test_file_path, type)
                else:
                    if type in src_type_codes or "type size_t" in type:
                        comment_code(test_file_path, type)

def remove_c_label_blocks(proj_path):
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")

    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            
            res = extract_c_label_blocks(file_path)
            c_label_blocks = res if res else []

            with open(file_path,'r') as file:
                content = file.read()
            for c_label_block in c_label_blocks:
                # print(c_label_block)
                content = content.replace(c_label_block, "")
            with open(file_path,'w') as file:
                file.write(content)

# 预处理C代码，移除Atomic类型
def remove_atomic_type(c_test_path):
    for root, dirs, files in os.walk(c_test_path):
        for file in files:
            file_path = os.path.join(root, file)

            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            
            content = content.replace("_Atomic", "")

            with open(file_path, 'w', encoding='utf-8') as file:
                file.write(content)

def add_alloc_set_limit(proj_path):
    line = "\n\talloc_test_set_limit(-1);"
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            functions = split_rust_functions_in_file(file_path)
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            for func in functions:
                if "test_" in func['name']:
                    origin_code = func['code']
                    pos = origin_code.find("{")
                    new_code = origin_code[:pos + 1] + line + origin_code[pos + 1:]
                    content = content.replace(origin_code, new_code)
            with open(file_path ,'w', encoding='utf-8') as file:
                file.write(content)

def add_print_num_assert_fun(proj_path):
    func_content = """
#[dtor]
fn print_num_assert() {
    let mut num_assert_guard = num_assert.lock().unwrap();
    println!("num_assert: {}", *num_assert_guard);
}
"""
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)

            with open(file_path, 'a', encoding='utf-8') as file:
                file.write(func_content)

def remove_no_mangle_annotation(proj_path):
    annotation = "#[no_mangle]"
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    SRC_DIR = os.path.join(PROJ_DIR, "src")
    for root, dirs, files in os.walk(SRC_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            content = content.replace(f"{annotation}\n", "")
            content = content.replace(annotation, "")
            with open(file_path, 'w', encoding='utf-8') as file:
                file.write(content)

    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            content = content.replace(f"{annotation}\n", "")
            content = content.replace(annotation, "")
            with open(file_path, 'w', encoding='utf-8') as file:
                file.write(content)

def test_static(proj_path):
    PROJ_DIR = proj_path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    SRC_DIR = os.path.join(PROJ_DIR, "src")
    for root, dirs, files in os.walk(SRC_DIR):
        for file in files:
            file_path = os.path.join(root, file)

            res = extract_static_declaration(file_path)
            static_code = '\n'.join(res) if res else ""

            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            if static_code not in content:
                print(f"file {file} not satisfy")
                print(static_code)
            # break


# print_functions="""
# #[dtor]
# fn print_num_assert() {
#     println!("num_assert: {}", num_assert.load(Ordering::Relaxed))
# }
# """
# def add_print_functions(proj_path):
#     PROJ_DIR = proj_path
#     TEST_DIR = os.path.join(PROJ_DIR, "tests")

#     for root, dirs, files in os.walk(TEST_DIR):
#         for file in files:
#             file_path = os.path.join(root, file)

#             with open(file_path, 'a') as file:
#                 file.write(print_functions)


if __name__ == "__main__":
    # add_crate_and_lib("../Output/primary")
    # remove_struct_and_type_declaration(f"Output/{CRATE_NAME}")
    # remove_extern_declaration_for_testfile(f"Output/{CRATE_NAME}")
    remove_no_mangle_annotation(f"Output/{CRATE_NAME}")
    # test_static(f"Output/{CRATE_NAME}")